import pyigl as igl
import numpy as np
import struct
from scipy import spatial
from iglhelpers import e2p, p2e
from PointCloud import PointCloud
import pickle
import time

data_path = '../../data/Cow/'
save_path = '../data/Cow/'

V = igl.eigen.MatrixXd()
F = igl.eigen.MatrixXi()
S = igl.eigen.MatrixXd()
I = igl.eigen.MatrixXi()
C = igl.eigen.MatrixXd()
N = igl.eigen.MatrixXd()

for frame in range(1,313):

    time_start = time.time()

    #读入obj mesh
    obj_file_name = "pos_%05d.obj"%(frame)
    igl.readOBJ(data_path + obj_file_name, V, F)

    txt_file_name = "pos_%d.txt"%(frame)
    f = open(data_path + txt_file_name,'rb')
    particle_num = struct.unpack('i',f.read(4))[0]
    query_data = np.zeros((particle_num,3),dtype=np.float32)
    mat_data = np.zeros((particle_num,4),dtype=np.float32)

    for i in range(particle_num):
        query_data[i][0] = mat_data[i][0] = struct.unpack('f',f.read(4))[0]
        query_data[i][1] = mat_data[i][1] = struct.unpack('f',f.read(4))[0]
        query_data[i][2] = mat_data[i][2] = struct.unpack('f',f.read(4))[0]

    f.close()

    P = p2e(query_data)

    print('begin computing signed distance')
    igl.signed_distance(P,V,F,igl.SignedDistanceType.SIGNED_DISTANCE_TYPE_PSEUDONORMAL,S,I,C,N)
    sd_value = e2p(S)

    print('begin neighbor querying')
    tree = spatial.KDTree(query_data)
    ball_radius = 0.01
    neighbor_lists = tree.query_ball_tree(other=tree, r=ball_radius)

    print('construct final point cloud data and label')
    pc_list = []
    for i in range(particle_num):
        pc = PointCloud(len(neighbor_lists[i]))
        center_point = query_data[i]
        for j in range(pc.size):
            pc.data[j] = (query_data[neighbor_lists[i][j]]-center_point)/ball_radius
            pc.label = sd_value[i][0]/ball_radius
        pc_list.append(pc)

    print('pickle dump data lists')
    pkl_file_name = "%d.pkl"%(frame)
    with open(save_path+pkl_file_name,'wb') as f:
        pickle.dump(pc_list, f)

    time_end = time.time()
    print('frame %d time cost %fs'%(frame,time_end-time_start))

